package com.learn.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Hex;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.AlgorithmParameterSpec;

/**
 * AES/GCM/NoPadding工具类
 *
 * @author yaoqihui
 * @version 2022/9/26 0026
 * @see [相关类/方法]
 * @since [ai-file-processor]
 */
@Slf4j
public class AESUtil
{
	private AESUtil ()
	{
	}

	private static final String AES = "AES";
	private static final String ALGORITHM = "AES/GCM/NoPadding";
	private static final Integer IV_SIZE = 16;

	/**
	 * AES/GCM/NoPadding 加密
	 *
	 * @param data    待加密的数据
	 * @param workKey 工作密钥
	 * @return 加密后的字符串
	 */
	public static String encrypt (String data, String workKey)
	{
		try
		{
			SecureRandom sr = SecureRandom.getInstance ("SHA1PRNG");
			byte[] iv = new byte[IV_SIZE];
			sr.nextBytes (iv);

			Cipher cipher = Cipher.getInstance (ALGORITHM);
			SecretKey secretKey = new SecretKeySpec (workKey.getBytes (), AES);
			//128 bit auth tag length
			AlgorithmParameterSpec parameterSpec = new GCMParameterSpec (128, iv);
			cipher.init (Cipher.ENCRYPT_MODE, secretKey, parameterSpec);

			byte[] cipherText = cipher.doFinal (data.getBytes (StandardCharsets.UTF_8));
			ByteBuffer byteBuffer = ByteBuffer.allocate (iv.length + cipherText.length);
			byteBuffer.put (iv);
			byteBuffer.put (cipherText);
			return Hex.encodeHexString (byteBuffer.array ());
		}
		catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException |
			   InvalidAlgorithmParameterException | IllegalBlockSizeException | BadPaddingException e)
		{
			log.error ("AES failed to encrypt: ", e);
		}
		return null;
	}

	/**
	 * AES/GCM/NoPadding解密
	 *
	 * @param encrypted 加密后的密码敏文
	 * @param workKey   工作秘钥
	 * @return 解密后的字节数组
	 */
	public static String decryptEnd (String encrypted, String workKey)
	{
		//截取iv向量
		String ivStr = encrypted.substring (0, 32);
		//截取真正的密文
		String encodeSource = encrypted.substring (32);
		byte[] result;
		try
		{
			result = decrypt (encodeSource, Hex.decodeHex (ivStr), workKey, ALGORITHM);
			if (result == null)
			{
				return null;
			}
		}
		catch (DecoderException e)
		{
			log.error ("AES failed to decryptEnd: ", e);
			return null;
		}
		return new String (result, StandardCharsets.UTF_8);
	}

	/**
	 * @param source    密文
	 * @param iv        vi12进制数组
	 * @param workKey   工作秘钥
	 * @param algorithm 算法 如：AES/GCM/NoPadding
	 * @return 解密后的字节数组
	 */
	public static byte[] decrypt (String source, byte[] iv, String workKey, String algorithm)
	{
		try
		{
			SecretKeySpec skey = new SecretKeySpec (workKey.getBytes (), AES);
			Cipher cipher = Cipher.getInstance (algorithm);
			GCMParameterSpec ivP = new GCMParameterSpec (128, iv);
			cipher.init (Cipher.DECRYPT_MODE, skey, ivP);
			return cipher.doFinal (Hex.decodeHex (source));
		}
		catch (BadPaddingException | IllegalBlockSizeException | NoSuchAlgorithmException | NoSuchPaddingException |
			   InvalidAlgorithmParameterException | InvalidKeyException | DecoderException e)
		{
			log.error ("AES failed to decrypt: ", e);
		}
		return null;
	}
}
